/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __LOGISTIC_LOSS_HH__
#define __LOGISTIC_LOSS_HH__

#include "dormandprince.hh"
#include "dyadicp.hh"
#include "loss.hh"
#include "wexpxmx.hh"

#include <algorithm>
#include <boost/optional.hpp>
#include <cmath>
#include <stdint.h>

#include "frsqrt.hh"

namespace flassol
{
  class LogisticLoss : public Loss
    {
      private:
        static LossFactory factory;

        uint64_t bad_labels;
        uint64_t npos;
        uint64_t nneg;

        uint64_t skipped_update;
        uint64_t n_smalls;
        uint64_t total_updates;

        class DFunc 
          {
            private:
              float adotb;
              float anormplusbnorm;
              float wdotx;
              float xnorm;
              float lambda;
              float y;

            public:
              DFunc (float _adotb,
                     float _anormplusbnorm,
                     float _wdotx,
                     float _xnorm,
                     float _lambda,
                     float _y)
                : adotb (_adotb),
                  anormplusbnorm (_anormplusbnorm),
                  wdotx (_wdotx),
                  xnorm (_xnorm),
                  lambda (_lambda),
                  y (_y)
                {
                }

              float
              operator() (float h,
                          float s) const
                {
                  float kappa = (fabsf (s) > 0) ? -lambda * h / s : 0.0f;

                  // TODO: rewrite dyadic::p in terms of decay not kappa

                  std::pair<bool, float> p = dyadic::p (s,
                                                        adotb,
                                                        anormplusbnorm,
                                                        wdotx,
                                                        xnorm,
                                                        kappa);

                  if (p.second * y > 10.0f)
                    {
                      return -y * expf (-y * p.second);
                    }
                  else
                    {
                      return -y / (1.0 + expf (y * p.second));
                    }
                }
          };

      public:
        explicit
        LogisticLoss () : bad_labels (0),
                          npos (0),
                          nneg (0),
                          skipped_update (0),
                          n_smalls (0),
                          total_updates (0)
          { 
          }

        virtual float 
        loss (float p,
              float z)
          {
            if (z != 1 && z != 0)
              {
                ++bad_labels;
              }

            float y = (z > 0) ? 1 : -1;

            if (y == 1) 
              { 
                ++npos; 
              } 
            else
              {
                ++nneg;
              }

            if (y * p > 10.0f)
              {
                // Log[1 + Exp[-x]] = Exp[-x] + O (Exp[-2x])
                
                return expf (-y * p);
              }
            else if (y * p < -10.0f)
              {
                // Log[1 + Exp[-x]] = Log[Exp[-x]] + Log[Exp[x] + 1]
                //                 = -x + O (Exp[x])
                
                return -y * p;
              }
            else
              {
                return logf (1.0f + expf (-y * p));
              }
          }

        virtual float 
        dldp (float p,
              float z)
          {
            float y = (z > 0) ? 1 : -1;

            if (y * p > 10.0f)
              {
                return -y * expf (-y * p);
              }
            else
              {
                return -y / (1.0 + expf (y * p));
              }
          }

        virtual std::pair<float, float>
        invariant_update (float p,
                          float wdotx,
                          float adotb,
                          float z,
                          float eta,
                          float anormplusbnorm,
                          float xnorm,
                          bool  adaptive,
                          float adaptivexnorm,
                          float lambda)
          {
            float y = (z > 0) ? 1 : -1;

            ++total_updates;

            if (anormplusbnorm == 0.0f)
              {
                float etaxtx = (adaptive) ? eta * adaptivexnorm * rsqrtf (xnorm)
                                          : eta;

                float xtx = (adaptive) ? adaptivexnorm : xnorm;

                // (W (exp (etaxtx + y p + exp (y p))) - etaxtx - exp (y p)) / (y xtx)
                // = (G (etaxtx + y p + exp (y p)) + y p) / (y xtx)
                // = (y G (etaxtx + y p + exp (y p)) + p) / xtx
                //
                // where G (x) = W (exp (x)) - x
                //
                // for very large x, G (x) = (-1 + 1/x) log (x) + O (1/x^2)
                //
                // so y G (exp (y p)) + p = y (-1 + exp (-y p)) y p + p + ...
                //                        = -y^2 p + p exp (-y p) + p
                //                        = p exp (-y p)

                return std::make_pair (
                         ( (y * p > 10.0f) 
                             ? p * expf (-y * p)
                             : y * wexpxmx (etaxtx + y * p + expf (y * p)) + p
                         )
                         / xtx,
                         1.0f);
              }
            else
              {
                // TODO: for quantile and hinge loss, the dynamics
                // stop when you hit the label.  this includes regularization,
                // so the user can specify a massive importance weight and
                // not kill the dyadic parameters.
                //
                // unfortunately here that is not the case: a huge learning
                // rate will not cause s(h) to increase much, but it will 
                // cause the decay factor to increase greatly

                // dyadic case
                
                float actualeta = (adaptive) ? eta * rsqrtf (xnorm)
                                             : eta / xnorm;
                float xtx = (adaptive) ? adaptivexnorm : xnorm;
                float norm = (anormplusbnorm + xtx);

                // first try small s approximation

                float p0y = (adotb + wdotx) * y;
                float smalls;

                if (p0y > 10.0f)
                  {
                    smalls = -actualeta * expf (-p0y) * y;
                  }
                else if (p0y < -10.0f)
                  {
                    smalls = -actualeta * y;
                  }
                else
                  {
                    smalls = -y * 
                              (1.0f + expf (-p0y)) * 
                              (1.0f - expf (-actualeta * norm / 
                                            (2.0f + 2.0f * coshf (p0y))))
                              / norm;
                  }

                if (fabsf (smalls) < 0.001f)
                  {
                    ++n_smalls;

                    return std::make_pair (smalls, 1.0f);
                  }
                else
                  {
                    //
                    // 
                    // s'[h] == -y / (1 + exp (y p))

                    DFunc dfunc (adotb, anormplusbnorm, wdotx, xtx, lambda, y);

                    boost::optional<float> s = 
                      dormand_prince_4_5 (dfunc,
                                          0.0f,
                                          0.0f,
                                          actualeta,
                                          1e-2f,
                                          actualeta / 1000.0f,
                                          actualeta,
                                          100);

                    if (s)
                      {
                        return std::make_pair (*s, expf (-actualeta * lambda));
                      }
                    else
                      {
                        ++skipped_update;

                        return std::make_pair (0.0f, 1.0f);
                      }
                  }
              }
          }

        virtual std::ostream& 
        report (std::ostream& o,
                LogLevel      level)
          {
            if (bad_labels > 0 && level <= WARN)
              {
                o << "WARNING: encountered " << bad_labels 
                  << " labels that were not 1 or 0"
                  << std::endl;
                o << "WARNING: positives were mapped to 1"
                  << " non-positives were mapped to 0"
                  << std::endl;
              }

            if (npos + nneg && level <= INFO)
              {
                double p =   static_cast<double> (npos) 
                           / static_cast<double> (npos + nneg);
                double Hp = (npos) ? npos * log (p) : 0;

                Hp += (nneg) ? nneg * log (1.0 - p) : 0;

                o << "unconditional probability of class 1 = " << p << std::endl;
                o << "best constant = ";

                if (p == 0)
                  {
                    o << "-infinity";
                  }
                else if (p == 1)
                  {
                    o << "+infinity";
                  }
                else
                  {
                    o << log (p / (1.0 - p));
                  }

                o << std::endl;

                o << "best constant loss = " << -Hp / (npos + nneg) << std::endl;
              }

            if (n_smalls > 0 && level <= WARN)
              {
                o << "WARNING: small s approximation utilized " 
                  << n_smalls 
                  << " times (" 
                  << ceil
                       (1000.0 * static_cast<double> (n_smalls) / 
                        static_cast<double> (total_updates)) / 10.0
                  << " %)"
                  << std::endl;
              }

            if (skipped_update > 0 && level <= WARN)
              {
                o << "WARNING: skipped " << skipped_update
                  << " updates ("
                  << ceil
                       (1000.0 * static_cast<double> (skipped_update) / 
                        static_cast<double> (total_updates)) / 10.0
                  << " %)"
                  << " due to loss of precision "
                  << std::endl;
              }

            return o;
          }
    };
}

#endif // __LOGISTIC_LOSS_HH__
